#ifndef AVLTREE_H
#define AVLTREE_H

#include <iomanip>
#include <iostream>
using namespace std;

template<class T>
class AVLTreeNode{
public:
    T m_key;
    AVLTreeNode *m_left;
    AVLTreeNode *m_right;
    int m_nHeight;

   AVLTreeNode(T value, AVLTreeNode *l, AVLTreeNode *r):
       m_key(value), m_left(l), m_right(r), m_nHeight(0) {}
};

template<class T>
class AVLTree
{
public:
    AVLTree();
    ~AVLTree();

    // 获取树的高度
    int height();
    // 比较两个值的大小
    int max(int a, int b);
    // 前序遍历
    void preOrder();
    // 中序遍历
    void inOrder();
    // 后序遍历
    void postOrder();

    // (递归实现)查找键值为key的节点
    AVLTreeNode<T>* search(T key);
    // (非递归实现)查找键值为key的节点
    AVLTreeNode<T>* iterativeSearch(T key);

    // 查找最小结点：返回最小结点的键值。
    T minimum();
    // 查找最大结点：返回最大结点的键值。
    T maximum();

    // 将结点(key为节点键值)插入到树中
    void insert(T key);

    // 删除结点(key为节点键值)
    void remove(T key);

    // 销毁二叉树
    void destroy();

    // 打印二叉树
    void print();

private:
    // 获取树的高度
    int height(AVLTreeNode<T>* tree) ;
    // 前序遍历"二叉树"
    void preOrder(AVLTreeNode<T>* tree);
    // 中序遍历"二叉树"
    void inOrder(AVLTreeNode<T>* tree);
    // 后序遍历"二叉树"
    void postOrder(AVLTreeNode<T>* tree);

    // (递归实现)查找键值为key的节点
    AVLTreeNode<T>* search(AVLTreeNode<T>* x, T key);
    // (非递归实现)查找中键值为key的节点
    AVLTreeNode<T>* iterativeSearch(AVLTreeNode<T>* x, T key);

    // 查找最小结点
    AVLTreeNode<T>* minimum(AVLTreeNode<T>* tree);
    // 查找最大结点
    AVLTreeNode<T>* maximum(AVLTreeNode<T>* tree);

    // LL：左左对应的情况(左单旋转)。
    AVLTreeNode<T>* leftLeftRotation(AVLTreeNode<T>* overbalance_node);

    // RR：右右对应的情况(右单旋转)。
    AVLTreeNode<T>* rightRightRotation(AVLTreeNode<T>* overbalance_node);

    // LR：左右对应的情况(左双旋转)。
    AVLTreeNode<T>* leftRightRotation(AVLTreeNode<T>* overbalance_node);

    // RL：右左对应的情况(右双旋转)。
    AVLTreeNode<T>* rightLeftRotation(AVLTreeNode<T>* overbalance_node);

    // 将结点(z)插入到树(tree)中
    AVLTreeNode<T> * insert(AVLTreeNode<T>* &tree, T key);

    // 删除树(tree)中的结点(z)，并返回被删除的结点
    AVLTreeNode<T>* remove(AVLTreeNode<T>* &tree, AVLTreeNode<T> *removeNode);

    // 销毁二叉树
    void destroy(AVLTreeNode<T>* &tree);

    // 打印二叉树
    void print(AVLTreeNode<T>* tree, T key, int direction);

private:
    AVLTreeNode<T> *m_root;
};

template<class T>
AVLTree<T>::AVLTree() : m_root(NULL)
{

}

template<class T>
AVLTree<T>::~AVLTree()
{
    destroy();
}

template<class T>
int AVLTree<T>::height(AVLTreeNode<T>* tree)
{
    if (tree != NULL)
        return tree->m_nHeight;

    return 0;
}

template<class T>
int AVLTree<T>::height()
{
    return height(m_root);
}

template<class T>
int AVLTree<T>::max(int a, int b)
{
    return a > b ? a : b;
}

template<class T>
void AVLTree<T>::preOrder(AVLTreeNode<T> * tree)
{
    if(tree != NULL)
    {
        cout << tree->m_key << " ";
        preOrder(tree->m_left);
        preOrder(tree->m_right);
    }
}

template<class T>
void AVLTree<T>::preOrder()
{
    preOrder(m_root);
}

template <class T>
void AVLTree<T>::inOrder(AVLTreeNode<T>* tree)
{
    if(tree != NULL)
    {
        inOrder(tree->m_left);
        cout << tree->m_key << " ";
        inOrder(tree->m_right);
    }
}

template<class T>
void AVLTree<T>::inOrder()
{
    inOrder(m_root);
}

template <class T>
void AVLTree<T>::postOrder(AVLTreeNode<T>* tree)
{
    if(tree != NULL)
    {
        postOrder(tree->m_left);
        postOrder(tree->m_right);
        cout << tree->m_key << " ";
    }
}

template<class T>
void AVLTree<T>::postOrder()
{
    postOrder(m_root);
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::search(AVLTreeNode<T>* x, T key)
{
    if(x == NULL || x->m_key == key)
        return x;

    if(x->m_key < key)
        return search(x->m_right, key);
    else
        return search(x->m_left, key);
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::search( T key)
{
    return search(m_root, key);
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::iterativeSearch(AVLTreeNode<T>* x, T key)
{
   while(x != NULL && x->m_key != key)
   {
       if( key < x->m_key)
           x = x->m_left;
       else
           x = x->m_right;
   }

   return x;
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::iterativeSearch( T key)
{
    return iterativeSearch(m_root, key);
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::minimum(AVLTreeNode<T>* tree)
{
    if(tree == NULL)
        return tree;

    while(tree->m_left != NULL)
        tree = tree->m_left;

    return tree;
}

template<class T>
T AVLTree<T>::minimum()
{
    return minimum(m_root)->m_key;
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::maximum(AVLTreeNode<T>* tree)
{
    if(tree == NULL)
        return tree;

    while(tree->m_right != NULL)
        tree = tree->m_right;

    return tree;
}

template<class T>
T AVLTree<T>::maximum()
{
    return maximum(m_root)->m_key;
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::leftLeftRotation(AVLTreeNode<T> *overbalance_node)
{
    auto left = overbalance_node->m_left;
    overbalance_node->m_left = left->m_right;
    left->m_right = overbalance_node;

    overbalance_node->m_nHeight = max(height(overbalance_node->m_right), height(overbalance_node->m_left)) + 1;
    left->m_nHeight = max(height(left->m_left), overbalance_node->m_nHeight) + 1;

    return left;
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::rightRightRotation(AVLTreeNode<T> *overbalance_node)
{
    auto right = overbalance_node->m_right;
    overbalance_node->m_right = right->m_left;
    right->m_left = overbalance_node;

    overbalance_node->m_nHeight = max(height(overbalance_node->m_left), height(overbalance_node->m_right)) + 1;
    right->m_nHeight = max(height(right->m_right), overbalance_node->m_nHeight) + 1;

    return right;
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::leftRightRotation(AVLTreeNode<T> *overbalance_node)
{
    overbalance_node->m_left = rightRightRotation(overbalance_node->m_left);

    return leftLeftRotation(overbalance_node);
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::rightLeftRotation(AVLTreeNode<T> *overbalance_node)
{
    overbalance_node->m_right = leftLeftRotation(overbalance_node->m_right);

    return rightRightRotation(overbalance_node);
}

template <class T>
AVLTreeNode<T> *AVLTree<T>::insert(AVLTreeNode<T> *&tree, T key)
{
    if(tree == NULL)
    {
        tree = new AVLTreeNode<T>(key, NULL, NULL);
        if(tree == NULL)
        {
            cout << " error" << endl;
        }
    }
    else if(key < tree->m_key)
    {
        tree->m_left = insert(tree->m_left, key);
        if(height(tree->m_left) - height(tree->m_right) >= 2)
        {
            if(key < tree->m_left->m_key)
                tree = leftLeftRotation(tree);
            else
                tree = leftRightRotation(tree);
        }
    }
    else if(key > tree->m_key)
    {
        tree->m_right = insert(tree->m_right, key);
        if (height(tree->m_right) - height(tree->m_left) == 2)
        {
            if (key > tree->m_right->m_key)
                tree = rightRightRotation(tree);
            else
                tree = rightLeftRotation(tree);
        }
    }

    tree->m_nHeight = max( height(tree->m_left), height(tree->m_right)) + 1;

    return tree;
}

template<class T>
void AVLTree<T>::insert(T key)
{
    insert(m_root, key);
}

template<class T>
AVLTreeNode<T>* AVLTree<T>::remove(AVLTreeNode<T> *&tree, AVLTreeNode<T> *removeNode)
{
    if( tree == NULL || removeNode == NULL)
        return NULL;

    if(removeNode->m_key < tree->m_key)
    {
        tree->m_left = remove(tree->m_left, removeNode);

        if(height(tree->m_right) - height(tree->m_left) >= 2)
        {
            AVLTreeNode<T> *overbalance_node = tree->m_right;
            if(height(overbalance_node->m_left) > height(overbalance_node->m_right))
                tree = rightLeftRotation(tree);
            else
                tree = rightRightRotation(tree);
        }
    }
    else if( removeNode->m_key > tree->m_key)
    {
        tree->m_right = remove(tree->m_right, removeNode);

        if(height(tree->m_left) - height(tree->m_right) >= 2)
        {
            AVLTreeNode<T> *overbalance_node = tree->m_left;

            if(height(overbalance_node->m_left) > height(overbalance_node->m_right))
                tree = leftLeftRotation(tree);
            else
                tree = leftRightRotation(tree);
        }
    }
    else
    {
        if(tree->m_left != NULL && tree->m_right != NULL)
        {
            if(height(tree->m_left) > height(tree->m_right))
            {
                AVLTreeNode<T> *max = maximum(tree->m_left);
                tree->m_key = max->m_key;
                tree->m_left = remove(tree->m_left, max);
            }
            else
            {
                AVLTreeNode<T> *min = minimum(tree->m_right);
                tree->m_key = min->m_key;
                tree->m_right = remove(tree->m_right, min);
            }
        }
        else
        {
             AVLTreeNode<T>* tmp = tree;
             tree = (tree->m_left!=NULL) ? tree->m_left : tree->m_right;
             delete tmp;
        }
    }

    return tree;
}

template <class T>
void AVLTree<T>::remove(T key)
{
    AVLTreeNode<T>* z;
    if ((z = search(m_root, key)) != NULL)
        m_root = remove(m_root, z);
}

template <class T>
void AVLTree<T>::destroy(AVLTreeNode<T>* &tree)
{
    if (tree==NULL)
        return ;

    if (tree->m_left != NULL)
        destroy(tree->m_left);
    else
        destroy(tree->m_right);

    delete tree;
}

template <class T>
void AVLTree<T>::destroy()
{
    destroy(m_root);
}

template <class T>
void AVLTree<T>::print(AVLTreeNode<T>* tree, T key, int direction)
{
    if(tree != NULL)
    {
        if(direction==0)
            cout << setw(2) << tree->m_key << " is root" << endl;
        else
            cout << setw(2) << tree->m_key << " is " << setw(2) << key << "'s "  << setw(12) << (direction==1?"right child" : "left child") << endl;

        print(tree->m_left, tree->m_key, -1);
        print(tree->m_right,tree->m_key,  1);
    }
}

template <class T>
void AVLTree<T>::print()
{
    if (m_root != NULL)
        print(m_root, m_root->m_key, 0);
}

#endif // AVLTREE_H
